#include "stdafx.h"
#include "vcclr.h"
#include "AutoArrayPtr.h"
#include "AsfFileWriter.h"

using namespace System;

namespace ManagedVC1
{
    using namespace Exceptions;

    AsfFileWriter::AsfFileWriter()
    {
        m_WMWriter = nullptr;
        m_WMWriterAdvanced = nullptr;
        m_WMProfileManager = nullptr;
        m_Profile = nullptr;
        m_currentInputNumber = 0;
        VideoSize = System::Drawing::Size::Empty;
    }

    void AsfFileWriter::Initialize(String^ fileName, 
        String^ profileName, 
        BYTE* privCodecData, 
        unsigned long privCodecDataLen,
        WAVEFORMATEX* pwfx
        )
    {
        if (String::IsNullOrEmpty(fileName)) throw gcnew ArgumentNullException();
        long hr = 0, width = 0, height = 0;

        CreateWMFObjects();

        pin_ptr<const wchar_t> wfileName = PtrToStringChars(fileName);
        hr = m_WMWriter->SetOutputFilename(wfileName);
        ExceptionUtility::ThrowExceptionForHr(hr, "IWMWriter::SetOutputFilename");

        // TODO update arguments and check here
        if (!String::IsNullOrEmpty(profileName) && pwfx != nullptr)
        {
            m_audioPresent = true;
            m_audioInputNumber = m_currentInputNumber++;

            if (profileName->Equals("default", StringComparison::OrdinalIgnoreCase))
            {
                if (m_Profile == nullptr)
                {
                    CreateEmptyProfile();
                }
                AddAudioStream(pwfx);
            }
            else
            {
                IWMProfile* pProf = nullptr;
                LoadCustomProfile(profileName, &pProf);
                m_Profile = pProf;
            }
        }

        if (!VideoSize->IsEmpty)
        {
            m_videoPresent = true;
            m_videoInputNumber = m_currentInputNumber++;

            AddVideoStream(privCodecData, privCodecDataLen);
        }

        hr = m_WMWriter->SetProfile(m_Profile);
        ExceptionUtility::ThrowExceptionForHr(hr, "IWMWriter::SetProfile");

        if (m_videoPresent)
        {
            hr = m_WMWriter->SetInputProps(m_videoInputNumber, nullptr);
            ExceptionUtility::ThrowExceptionForHr(hr, "IWMWriter::SetInputProps");
        }

        SetCodecInfo();

        hr = m_WMWriter->BeginWriting();
        ExceptionUtility::ThrowExceptionForHr(hr, "IWMWriter::BeginWriting");
    }

    void AsfFileWriter::CreateWMFObjects()
    {
        long hr = 0;

        IWMWriter* pWriter = nullptr;
        hr = WMCreateWriter(nullptr, &pWriter);
        ExceptionUtility::ThrowExceptionForHr(hr);
        m_WMWriter = pWriter;

        IWMWriterAdvanced* pWriterAdv = nullptr;
        hr = m_WMWriter->QueryInterface(IID_IWMWriterAdvanced, (void **)&pWriterAdv);
        ExceptionUtility::ThrowExceptionForHr(hr);
        m_WMWriterAdvanced = pWriterAdv;        

        IWMProfileManager* pProfileMgr = nullptr;
        hr = WMCreateProfileManager(&pProfileMgr);
        ExceptionUtility::ThrowExceptionForHr(hr);
        m_WMProfileManager = pProfileMgr;
    }

    void AsfFileWriter::SetCodecInfo()
    {
        long hr = 0;
        IWMHeaderInfo3* pHeaderInfo;

        hr = m_WMWriter->QueryInterface(IID_IWMHeaderInfo3, (void **)&pHeaderInfo);
        ExceptionUtility::ThrowExceptionForHr(hr);

        WMT_CODEC_INFO_TYPE codecType;
        AutoArrayPtr<BYTE> codecInfo(new BYTE[4]);

        pin_ptr<const wchar_t> vc1str = PtrToStringChars("VC-1");
        pin_ptr<const wchar_t> simpleStr = PtrToStringChars("VC-1 Simple Profile");
        pin_ptr<const wchar_t> mainStr = PtrToStringChars("VC-1 Main Profile");
        pin_ptr<const wchar_t> advStr = PtrToStringChars("VC-1 Advanced Profile");

        String^ typeString;
        WORD pwIndex = 0;

        switch (VC1Profile)
        {
        case 0:
            codecInfo[0] = 0x57;
            codecInfo[1] = 0x4D;
            codecInfo[2] = 0x56;
            codecInfo[3] = 0x33;

            codecType = WMT_CODECINFO_VIDEO;

            hr = pHeaderInfo->AddCodecInfo((LPWSTR)vc1str, (LPWSTR)simpleStr, codecType, 4, codecInfo);
            ExceptionUtility::ThrowExceptionForHr(hr);
            break;
        case 1:
            codecInfo[0] = 0x57;
            codecInfo[1] = 0x4D;
            codecInfo[2] = 0x56;
            codecInfo[3] = 0x33;

            codecType = WMT_CODECINFO_VIDEO;

            hr = pHeaderInfo->AddCodecInfo((LPWSTR)vc1str, (LPWSTR)mainStr, codecType, 4, codecInfo);
            ExceptionUtility::ThrowExceptionForHr(hr);
            typeString = "MP@ML";
            break;
        case 2:
            codecInfo[0] = 0x57;
            codecInfo[1] = 0x56;
            codecInfo[2] = 0x43;
            codecInfo[3] = 0x31;

            codecType = WMT_CODECINFO_VIDEO;

            hr = pHeaderInfo->AddCodecInfo((LPWSTR)vc1str, (LPWSTR)advStr, codecType, 4, codecInfo);
            ExceptionUtility::ThrowExceptionForHr(hr);
            typeString = "AP@L1";
            break;
        }

        if (!String::IsNullOrEmpty(typeString))
        {
            pin_ptr<const wchar_t> typeStr = PtrToStringChars(typeString);
            hr = pHeaderInfo->AddAttribute(VideoStreamNumber, g_wszDeviceConformanceTemplate, &pwIndex, WMT_TYPE_STRING, 0, (BYTE*)typeStr, sizeof(WCHAR) * typeString->Length);
            ExceptionUtility::ThrowExceptionForHr(hr);
        }

        SAFE_RELEASE(pHeaderInfo);
    }

    void AsfFileWriter::CreateEmptyProfile()
    {
        IWMProfile* pProf = nullptr;
        long hr = m_WMProfileManager->CreateEmptyProfile(WMT_VER_9_0, &pProf);
        ExceptionUtility::ThrowExceptionForHr(hr);
        if (pProf == nullptr) throw gcnew OutOfMemoryException();
        m_Profile = pProf;
    }

    void AsfFileWriter::WriteVideoFrame(BYTE* vc1Data, unsigned long vc1DataSize, __int64 timeStamp, FrameType ft)
    {
        if (!m_videoPresent) return;

        long hr = 0;
        INSSBuffer* pSample = nullptr;
        BYTE* pBuffer = nullptr;
        unsigned long flags = 0;

        if (ft == I || ft == II || ft == IP) flags = WM_SF_CLEANPOINT;

		try
		{
			hr = m_WMWriter->AllocateSample(vc1DataSize, &pSample);
			ExceptionUtility::ThrowExceptionForHr(hr);

			hr = pSample->GetBuffer(&pBuffer);
			ExceptionUtility::ThrowExceptionForHr(hr);

			hr = pSample->SetLength(vc1DataSize);
			ExceptionUtility::ThrowExceptionForHr(hr);

			memcpy_s(pBuffer, vc1DataSize, vc1Data, vc1DataSize);
			LONGLONG timePerFrame = ( (LONGLONG) 10000000 ) / ( FrameRate->Rate );
			hr = m_WMWriterAdvanced->WriteStreamSample(VideoStreamNumber, timeStamp, 0L, timePerFrame, flags, pSample);
			ExceptionUtility::ThrowExceptionForHr(hr);
		}
		finally
		{
			SAFE_RELEASE(pSample);
		}
    }

    void AsfFileWriter::WriteAudioFrame(BYTE* audioData, unsigned long audioSize, __int64 timeStamp)
    {
        if (!m_audioPresent) return;

        long hr = 0;
        INSSBuffer* pSample = nullptr;
        BYTE* pBuffer = nullptr;

        try
        {
            hr = m_WMWriter->AllocateSample(audioSize, &pSample);
            ExceptionUtility::ThrowExceptionForHr(hr, "IWMWriter::AllocateSample");

            hr = pSample->GetBuffer(&pBuffer);
            ExceptionUtility::ThrowExceptionForHr(hr, "INSSBuffer::GetBuffer");

            hr = pSample->SetLength(audioSize);
            ExceptionUtility::ThrowExceptionForHr(hr, "INSSBuffer::SetLength");

            memcpy_s(pBuffer, audioSize, audioData, audioSize);
            hr = m_WMWriter->WriteSample(m_audioInputNumber, timeStamp, 0, pSample);
            ExceptionUtility::ThrowExceptionForHr(hr, "IWMWriter::WriteSample");
        }
        finally
        {
            SAFE_RELEASE(pSample);
        }
    }

    void AsfFileWriter::LoadCustomProfile(String^ profileFile, IWMProfile** ppProfile)
    {
        String^ profileData = System::IO::File::ReadAllText(profileFile);
        pin_ptr<const wchar_t> pdata = PtrToStringChars(profileData);
        long hr = m_WMProfileManager->LoadProfileByData(pdata, ppProfile);
        ExceptionUtility::ThrowExceptionForHr(hr);
    }

    void AsfFileWriter::AddAudioStream(WAVEFORMATEX *pWaveLimits)
    {
        IWMStreamConfig* pConfig = nullptr;
        IWMMediaProps* pProps = nullptr;
        IWMCodecInfo3* pCodecInfo = nullptr;
        WAVEFORMATEX* pWave = nullptr;
        AutoArrayPtr<BYTE> pType;
        
        const unsigned long InvalidIndex = 0xFFFF;

        unsigned long bestRate = 0;
        unsigned long entries = 0;
        unsigned long typeLen = 0;
        unsigned long codecIndex = InvalidIndex;
        long hr = 0;

        hr = m_WMProfileManager->QueryInterface(IID_IWMCodecInfo, (void **)&pCodecInfo);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pCodecInfo->GetCodecInfoCount(WMMEDIATYPE_Audio, &entries);
        ExceptionUtility::ThrowExceptionForHr(hr);

        for (unsigned long index = 0; index < entries && codecIndex == InvalidIndex; index++)
        {
            try
            {
                hr = pCodecInfo->GetCodecFormat(WMMEDIATYPE_Audio, index, 0, &pConfig);
                ExceptionUtility::ThrowExceptionForHr(hr);

                hr = pConfig->QueryInterface(IID_IWMMediaProps, (void **)&pProps);
                ExceptionUtility::ThrowExceptionForHr(hr);

                hr = pProps->GetMediaType(nullptr, &typeLen);
                ExceptionUtility::ThrowExceptionForHr(hr);

                pType.Reset(new BYTE[typeLen]);
                if (pType.GetPointer() == nullptr) throw gcnew OutOfMemoryException();

                hr = pProps->GetMediaType((WM_MEDIA_TYPE*)pType.GetPointer(), &typeLen);
                ExceptionUtility::ThrowExceptionForHr(hr);

                WM_MEDIA_TYPE* mediaType = (WM_MEDIA_TYPE*)pType.GetPointer();
                if (mediaType->subtype == WMMEDIASUBTYPE_WMAudioV8) codecIndex = index;
            }
            finally
            {
                SAFE_RELEASE(pConfig);
                SAFE_RELEASE(pProps);
                pType.Reset();
            }
        }

        if (codecIndex == InvalidIndex) throw gcnew ArgumentException();

        hr = pCodecInfo->GetCodecFormatCount(WMMEDIATYPE_Audio, codecIndex, &entries);
        ExceptionUtility::ThrowExceptionForHr(hr);

        for (unsigned long index = 0; index < entries; index++)
        {
            hr = pCodecInfo->GetCodecFormat(WMMEDIATYPE_Audio, codecIndex, index, &pConfig);
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pConfig->QueryInterface(IID_IWMMediaProps, (void**)&pProps);
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pProps->GetMediaType(nullptr, &typeLen);
            ExceptionUtility::ThrowExceptionForHr(hr);

            pType.Reset(new BYTE[typeLen]);
            if (pType.GetPointer() == nullptr) throw gcnew OutOfMemoryException();

            hr = pProps->GetMediaType((WM_MEDIA_TYPE*)pType.GetPointer(), &typeLen);
            ExceptionUtility::ThrowExceptionForHr(hr);

            WM_MEDIA_TYPE* mediaType = (WM_MEDIA_TYPE*)pType.GetPointer();
            if (mediaType->cbFormat < sizeof(WAVEFORMATEX)) throw gcnew Exception("Expected a WAVEFORMATEX structure");

            pWave = (WAVEFORMATEX*)mediaType->pbFormat; 

            // Do not check particulars unless the bit rate is in range.
            if((pWave->nAvgBytesPerSec * 8) > bestRate &&
               (pWave->nAvgBytesPerSec * 8) <= Bitrate)
            {
                // Check the limits.
                if((pWave->nChannels == pWaveLimits->nChannels) &&
                   (pWave->wBitsPerSample == pWaveLimits->wBitsPerSample) &&
                   (pWave->nSamplesPerSec == pWaveLimits->nSamplesPerSec))
                {                    
                    if((pWave->nAvgBytesPerSec / pWave->nBlockAlign) >= 
                           ((pWave->nAvgBytesPerSec >= 4000) ? 5.0 : 3.0))
                    {
                        // Found a match. We'll use the first good match we find.                          

                        // Set the bit rate.
                        bestRate = (pWave->nAvgBytesPerSec * 8);

                        break;

                    }                   
                } // if matching limits
            } // if valid bit rate

            SAFE_RELEASE(pConfig);
            SAFE_RELEASE(pProps);
            pWave = nullptr;
            pType.Reset();
        }

        if (pConfig != nullptr)
        {
            SetStreamBasics(pConfig, bestRate, (WM_MEDIA_TYPE*)pType.GetPointer());
        }
        else throw gcnew HrException("Unable to find a suitable audio codec", NS_E_AUDIO_CODEC_NOT_INSTALLED);

        hr = m_Profile->AddStream(pConfig);
        ExceptionUtility::ThrowExceptionForHr(hr);

        SAFE_RELEASE(pCodecInfo);
        SAFE_RELEASE(pConfig);
        SAFE_RELEASE(pProps);
        pType.Reset();
    }

    void AsfFileWriter::SetStreamBasics(IWMStreamConfig* pIWMStreamConfig, DWORD dwBitrate, WM_MEDIA_TYPE* pmt)
    {
        if (pIWMStreamConfig == nullptr || pmt == nullptr) throw gcnew ArgumentNullException();

        long hr = 0;
        IWMMediaProps* pMediaProps = nullptr;

        try
        {
            hr = pIWMStreamConfig->SetStreamNumber(AudioStreamNumber);
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pIWMStreamConfig->SetBitrate(dwBitrate);
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pIWMStreamConfig->QueryInterface(IID_IWMMediaProps, (void**)&pMediaProps);
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pMediaProps->SetMediaType(pmt);
            ExceptionUtility::ThrowExceptionForHr(hr);
        }
        finally
        {
            SAFE_RELEASE(pMediaProps);
        }
    }

    void AsfFileWriter::AddVideoStream(BYTE* privCodecData, unsigned long privCodecDataLen)
    {
        long hr = 0;
        unsigned long cbMediaType = 0;
        IWMStreamConfig* pConfig = nullptr;
        IWMMediaProps* pMediaProps = nullptr;
        IWMPropertyVault *pVault = nullptr;
        WM_MEDIA_TYPE* pMediaType = nullptr;
        WMVIDEOINFOHEADER *pVIH = nullptr;

        if (m_Profile == nullptr) CreateEmptyProfile();

        hr = m_Profile->CreateNewStream(WMMEDIATYPE_Video, &pConfig);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pConfig->SetStreamNumber(VideoStreamNumber);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pConfig->SetConnectionName(L"SDKTest");
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = m_Profile->AddStream(pConfig);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pConfig->QueryInterface(IID_IWMMediaProps, (void **)&pMediaProps);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pMediaProps->GetMediaType(nullptr, &cbMediaType);
        ExceptionUtility::ThrowExceptionForHr(hr);

        pMediaType = (WM_MEDIA_TYPE*)new BYTE[cbMediaType + privCodecDataLen];
        if (pMediaType == nullptr) throw gcnew OutOfMemoryException();

        ZeroMemory(pMediaType, cbMediaType + privCodecDataLen);
        hr = pMediaProps->GetMediaType(pMediaType, &cbMediaType);
        ExceptionUtility::ThrowExceptionForHr(hr);

        pMediaType->bTemporalCompression = TRUE;

        pVIH = (WMVIDEOINFOHEADER*)pMediaType->pbFormat;

        pVIH->dwBitRate = (DWORD)Bitrate;
        pVIH->rcSource.left = 0;
        pVIH->rcSource.top = 0;
        pVIH->rcSource.right = VideoSize->Width;
        pVIH->rcSource.bottom = VideoSize->Height;
        pVIH->rcTarget.left = 0;
        pVIH->rcTarget.top = 0;
        pVIH->rcTarget.right = VideoSize->Width;
        pVIH->rcTarget.bottom = VideoSize->Height;
        pVIH->bmiHeader.biWidth = VideoSize->Width;
        pVIH->bmiHeader.biHeight = VideoSize->Height;
        pVIH->bmiHeader.biSizeImage = FrameSizeInBytes;
        pVIH->AvgTimePerFrame = ( (LONGLONG) 10000000 ) / ( FrameRate->Rate );

        // Configure members based on profile.
        switch(VC1Profile)
        {
            case 0:
            case 1:
                // Simple or Main.
                pVIH->bmiHeader.biCompression = FOURCC_WMV3;
                pMediaType->subtype = WMMEDIASUBTYPE_WMV3;
                break;
            case 2:
                // Advanced
                pVIH->bmiHeader.biCompression = FOURCC_WVC1;
                pMediaType->subtype = MEDIASUBTYPE_WVC1;
                break;
            default:
                // Unexpected.
                throw gcnew ArgumentException();
                break;
        }

        // Set the format size.
        pMediaType->cbFormat += privCodecDataLen;
       
        // Append the VC-1 header information to the end of the WM_MEDIA_TYPE
        // structure.
        memcpy((BYTE*)pMediaType + cbMediaType, privCodecData, privCodecDataLen);

        hr = pConfig->SetBitrate((DWORD)Bitrate);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pConfig->SetBufferWindow(BufferInMilliseconds);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pMediaProps->SetMediaType(pMediaType);
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = pConfig->QueryInterface(IID_IWMPropertyVault, (void**)&pVault);
        ExceptionUtility::ThrowExceptionForHr(hr);

        WMT_ATTR_DATATYPE wad = WMT_TYPE_BOOL;
        unsigned long temp = 0;

        BOOL vbr = IsVBR;
        hr = pVault->SetProperty(g_wszVBREnabled, wad, (BYTE*)&vbr, sizeof(WMT_TYPE_BOOL));
        ExceptionUtility::ThrowExceptionForHr(hr);

        wad = WMT_TYPE_DWORD;

        if (IsVBR)
        {
            hr = pVault->SetProperty(g_wszVBRBitrateMax, wad, (BYTE*)&temp, sizeof(WMT_TYPE_DWORD));
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pVault->SetProperty(g_wszVBRBufferWindowMax, wad, (BYTE*)&temp, sizeof(WMT_TYPE_DWORD));
            ExceptionUtility::ThrowExceptionForHr(hr);

            hr = pVault->SetProperty(g_wszVBRQuality, wad, (BYTE*)&temp, sizeof(WMT_TYPE_DWORD));
            ExceptionUtility::ThrowExceptionForHr(hr);

            double br = Bitrate;
            hr = pVault->SetProperty(L"_RAVG", wad, (BYTE*)&br, sizeof(WMT_TYPE_DWORD));
            ExceptionUtility::ThrowExceptionForHr(hr);
        }

        hr = m_Profile->SetName(L"ASF Writer");
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = m_Profile->SetDescription(L"Writes VC-1 compressed video and uncompressed audio into ASF file");
        ExceptionUtility::ThrowExceptionForHr(hr);

        hr = m_Profile->ReconfigStream(pConfig);
        ExceptionUtility::ThrowExceptionForHr(hr);

        SAFE_ARRAY_DELETE(pMediaType);
        SAFE_RELEASE(pConfig);
        SAFE_RELEASE(pVault);
        SAFE_RELEASE(pMediaProps);
    }

    void AsfFileWriter::Stop()
    {
        if (m_WMWriter != nullptr)
        {
            m_WMWriter->EndWriting();
        }
    }
}
